__kernel
void resetGrid(
	__global int* countGrid,
	__global int4* indexGrid
) {
	unsigned int flatGridIndex = get_global_id(0);

	countGrid[flatGridIndex] = 0;
	indexGrid[flatGridIndex].x = -1;
	indexGrid[flatGridIndex].y = -1;
	indexGrid[flatGridIndex].z = -1;
	indexGrid[flatGridIndex].w = -1;

}

__kernel
void updateGrid(
	__global int* countGrid,
	__global int4* indexGrid,
	__global float3* particlePosition,
	float3 minimumPosition,
	float voxelSideLength,
	int gridSideLength,
	__global int3* particleGridIndex
) {
	unsigned int particleIndex = get_global_id(0);

	particleGridIndex[particleIndex].x =  (int)((particlePosition[particleIndex].x - minimumPosition.x) / voxelSideLength);
	particleGridIndex[particleIndex].y =  (int)((particlePosition[particleIndex].y - minimumPosition.y) / voxelSideLength);
	particleGridIndex[particleIndex].z =  (int)((particlePosition[particleIndex].z - minimumPosition.z) / voxelSideLength);
	
	bool validIndex = (particleGridIndex[particleIndex].x > 0) && 
		(particleGridIndex[particleIndex].x < gridSideLength-1) &&
		(particleGridIndex[particleIndex].y > 0) &&
		(particleGridIndex[particleIndex].y < gridSideLength-1) &&
		(particleGridIndex[particleIndex].z > 0) &&
		(particleGridIndex[particleIndex].z < gridSideLength-1);

	if (validIndex) {
		int xStride = gridSideLength * gridSideLength;
		int yStride = gridSideLength;
		int flatGridIndex = particleGridIndex[particleIndex].x*xStride + 
			particleGridIndex[particleIndex].y * yStride + 
			particleGridIndex[particleIndex].z;

		int particlesInCell = atomic_inc(&countGrid[flatGridIndex]);

		if (particlesInCell == 3) {
			indexGrid[flatGridIndex].w = particleIndex;
		} else if (particlesInCell == 2) {
			indexGrid[flatGridIndex].z = particleIndex;
		} else if (particlesInCell == 1) {
			indexGrid[flatGridIndex].y = particleIndex;
		} else if (particlesInCell == 0) {
			indexGrid[flatGridIndex].x = particleIndex;
		}
	}
}


__kernel
void collisionDetection(
	__global float3* particleMass,
	__global float3* particlePosition,
	__global float3* particleVelocity,
	__global float3* particleForce,
	float particleRadius,
	float worldSize,
	float springCoefficient,
	float dampingCoefficient,
	__global int3* particleGridIndex,
	__global int* countGrid,
	__global int4* indexGrid,
	int gridSideLength
) {

	unsigned int particleIndex = get_global_id(0);

	particleForce[particleIndex].x = 0.0f;
	particleForce[particleIndex].y = 0.0f; 
	particleForce[particleIndex].z = 0.0f; 

	int3 gridIndex = particleGridIndex[particleIndex];

	//Pretend border cell is 1 position inwards to avoid checking outside bounds for neighbors
	gridIndex = clamp(gridIndex, 1, gridSideLength-2);

	int xStride = gridSideLength*gridSideLength;
	int yStride = gridSideLength;

	int flatGridIndex = gridIndex.x * xStride +
		gridIndex.y * yStride + gridIndex.z;

	int4 neighborCells[27];
	int cellIndexJ = 0;

	flatGridIndex -= xStride;
	flatGridIndex += 2*yStride;
	flatGridIndex += 2; //zStride


	for (int x=0; x<3; x++) {
		flatGridIndex -= 3*yStride;

		for (int y=0; y<3; y++) {
			flatGridIndex -= 3;

			for (int z=0; z<3; z++) {

				neighborCells[cellIndexJ] = indexGrid[flatGridIndex];
				cellIndexJ++;
				flatGridIndex++;
			}
			flatGridIndex += yStride;
		}
		flatGridIndex += xStride;
	}

	for (int j=0; j<27; j++) {
		int neighborParticles[4] = {
			neighborCells[j].x,
			neighborCells[j].y,
			neighborCells[j].z,
			neighborCells[j].w };

		for (int k=0; k<4; k++) {
			int otherParticle = neighborParticles[k];
			if ((otherParticle != particleIndex) && (otherParticle != (-1))) {
				float3 distance = particlePosition[otherParticle] - particlePosition[particleIndex];

				float absDistance = sqrt(distance.x*distance.x +
						distance.y*distance.y +
						distance.z*distance.z);

				if ((absDistance + 0.000001f) < (2.0f * particleRadius)) {
					particleForce[particleIndex].x -= springCoefficient*
						(particleRadius+particleRadius - absDistance)*(distance.x/absDistance);
					particleForce[particleIndex].y -= springCoefficient*
						(particleRadius+particleRadius - absDistance)*(distance.y/absDistance);
					particleForce[particleIndex].z -= springCoefficient*
						(particleRadius+particleRadius - absDistance)*(distance.z/absDistance);

					float3 relativeVelocity = {
						particleVelocity[otherParticle].x - particleVelocity[particleIndex].x,
						particleVelocity[otherParticle].y - particleVelocity[particleIndex].y,
						particleVelocity[otherParticle].z - particleVelocity[particleIndex].z };

					particleForce[particleIndex].x += dampingCoefficient*relativeVelocity.x;
					particleForce[particleIndex].y += dampingCoefficient*relativeVelocity.y;
					particleForce[particleIndex].z += dampingCoefficient*relativeVelocity.z;

				}
			}
		}
	}


	//Boundary forces
	{
		bool collisionOccured = false;
		// Ground collision
		if (particlePosition[particleIndex].y-particleRadius < 0.0f) {
			collisionOccured = true;
			particleForce[particleIndex].y += springCoefficient*
				(particleRadius-particlePosition[particleIndex].y);
		}

		// X-axis Wall Collision
		if (particlePosition[particleIndex].x-particleRadius < -worldSize) { 
			collisionOccured = true;
			particleForce[particleIndex].x += springCoefficient*
				(-worldSize - particlePosition[particleIndex].x + particleRadius);

		} else if (particlePosition[particleIndex].x+particleRadius > worldSize) {
			collisionOccured = true;
			particleForce[particleIndex].x += springCoefficient*
				(worldSize - particlePosition[particleIndex].x - particleRadius);
		}

		// Z-axis Wall Collision
		if (particlePosition[particleIndex].z-particleRadius < -worldSize) {
			collisionOccured = true;
			particleForce[particleIndex].z += springCoefficient*
				(-worldSize - particlePosition[particleIndex].z + particleRadius);

		} else if (particlePosition[particleIndex].z+particleRadius > worldSize) {
			collisionOccured = true;
			particleForce[particleIndex].z += springCoefficient*
				(worldSize - particlePosition[particleIndex].z - particleRadius);
		}

		// Damping
		if (collisionOccured) {
			particleForce[particleIndex].x -= dampingCoefficient*particleVelocity[particleIndex].x;
			particleForce[particleIndex].y -= dampingCoefficient*particleVelocity[particleIndex].y;
			particleForce[particleIndex].z -= dampingCoefficient*particleVelocity[particleIndex].z;
		}
	}
}


__kernel
void updateMomenta(
		__global float* mass,
		__global float3* force,
		__global float3* position,
		__global float3* linearMomentum,
		__global float3* angularMomentum,
		__global float3* particlePosition,
		__global float3* particleForce,
		float delta,
		float terminalVelocity
) {
	unsigned int bodyIndex = get_global_id(0);
	unsigned int totalNumberOfParticles = get_global_size(0) * 27;
	unsigned int particleIndex = bodyIndex * 27;

	force[bodyIndex].x = 0.0f;
	force[bodyIndex].y = mass[bodyIndex] * -9.81f; //force of gravity
	force[bodyIndex].z = 0.0f;

	float3 torque = {0.0f, 0.0f, 0.0f};

	//Calculate body force and torque
	for (int i=0; i<27; i++) {
		force[bodyIndex] += particleForce[particleIndex+i];
		float3 relativePos = particlePosition[particleIndex+i] - position[bodyIndex];
		torque += cross(relativePos, particleForce[particleIndex+i]);
	}

	float terminalMomentum = terminalVelocity * mass[bodyIndex];

	linearMomentum[bodyIndex].x += force[bodyIndex].x * delta;
	linearMomentum[bodyIndex].y += force[bodyIndex].y * delta;
	linearMomentum[bodyIndex].z += force[bodyIndex].z * delta;

	//Limit momentum by terminal momentum
	linearMomentum[bodyIndex] = clamp(linearMomentum[bodyIndex], -terminalMomentum, terminalMomentum);

	angularMomentum[bodyIndex].x += torque.x * delta;
	angularMomentum[bodyIndex].y += torque.y * delta;
	angularMomentum[bodyIndex].z += torque.z * delta;
}

__kernel
void performStep(
	__global float* mass,
	__global float3* position,
	__global float3* velocity,
	__global float3* linearMomentum,
	__global float4* quaternion,
	__global float* rotationMatrix,
	__global float3* angularVelocity,
	__global float3* angularMomentum,
	__global float3* initialIITDiagonal,
	__global float* inverseInertiaTensor,
	__global float* bodyVBO,
	float delta,
	float particleRadius,
	int bodyVBOStride
) {
	unsigned int bodyIndex = get_global_id(0);
	unsigned int bodyVBOIndex = bodyIndex * 24 * 3;
	unsigned int matrixIndex = bodyIndex * 9;

	//Update inverse inertia tensor
	{
		float a = rotationMatrix[matrixIndex];
		float b = rotationMatrix[matrixIndex+1];
		float c = rotationMatrix[matrixIndex+2];
		float d = rotationMatrix[matrixIndex+3];
		float e = rotationMatrix[matrixIndex+4];
		float f = rotationMatrix[matrixIndex+5];
		float g = rotationMatrix[matrixIndex+6];
		float h = rotationMatrix[matrixIndex+7];
		float i = rotationMatrix[matrixIndex+8];

		float u = initialIITDiagonal[bodyIndex].x;
		float v = initialIITDiagonal[bodyIndex].y;
		float w = initialIITDiagonal[bodyIndex].z;

		inverseInertiaTensor[matrixIndex] = u*a*a + b*b*v + c*c*w;
		inverseInertiaTensor[matrixIndex+1] = a*d*u + b*e*v + c*f*w;
		inverseInertiaTensor[matrixIndex+2] = a*g*u + b*h*v + c*i*w;
		inverseInertiaTensor[matrixIndex+3] = a*d*u + b*e*v + c*f*w;
		inverseInertiaTensor[matrixIndex+4] = u*d*d + e*e*v + f*f*w;
		inverseInertiaTensor[matrixIndex+5] = d*g*u + e*h*v + f*i*w;
		inverseInertiaTensor[matrixIndex+6] = a*g*u + b*h*v + c*i*w;
		inverseInertiaTensor[matrixIndex+7] = d*g*u + e*h*v + f*i*w;
		inverseInertiaTensor[matrixIndex+8] = u*g*g + h*h*v + i*i*w;
	}

	//Perform linear step
	{
		velocity[bodyIndex].x = linearMomentum[bodyIndex].x / mass[bodyIndex];
		velocity[bodyIndex].y = linearMomentum[bodyIndex].y / mass[bodyIndex];
		velocity[bodyIndex].z = linearMomentum[bodyIndex].z / mass[bodyIndex];

		position[bodyIndex].x += velocity[bodyIndex].x * delta;
		position[bodyIndex].y += velocity[bodyIndex].y * delta;
		position[bodyIndex].z += velocity[bodyIndex].z * delta;
	}

	//Perform angular step
	{
		//Update angular velocity
		{
			float a = inverseInertiaTensor[matrixIndex];
			float b = inverseInertiaTensor[matrixIndex+1];
			float c = inverseInertiaTensor[matrixIndex+2];
			float d = inverseInertiaTensor[matrixIndex+3];
			float e = inverseInertiaTensor[matrixIndex+4];
			float f = inverseInertiaTensor[matrixIndex+5];
			float g = inverseInertiaTensor[matrixIndex+6];
			float h = inverseInertiaTensor[matrixIndex+7];
			float i = inverseInertiaTensor[matrixIndex+8];

			float u = angularMomentum[bodyIndex].x;
			float v = angularMomentum[bodyIndex].y;
			float w = angularMomentum[bodyIndex].z;

			angularVelocity[bodyIndex].x = a*u + b*v + c*w;
			angularVelocity[bodyIndex].y = d*u + e*v + f*w;
			angularVelocity[bodyIndex].z = g*u + h*v + i*w;
		}
		float angularVelocitySize = sqrt(
				angularVelocity[bodyIndex].x*angularVelocity[bodyIndex].x +
				angularVelocity[bodyIndex].y*angularVelocity[bodyIndex].y +
				angularVelocity[bodyIndex].z*angularVelocity[bodyIndex].z ); 

		if (angularVelocitySize > 0.0f) {
			float3 rotationAxis = {
				angularVelocity[bodyIndex].x / angularVelocitySize,
				angularVelocity[bodyIndex].y / angularVelocitySize,
				angularVelocity[bodyIndex].z / angularVelocitySize};

			float rotationAngle = angularVelocitySize*delta;

			float ds = cos(rotationAngle/2.0f);
			float dvx = rotationAxis.x*sin(rotationAngle/2.0f);
			float dvy = rotationAxis.y*sin(rotationAngle/2.0f);
			float dvz = rotationAxis.z*sin(rotationAngle/2.0f);

			float s = quaternion[bodyIndex].x;
			float vx = quaternion[bodyIndex].y;
			float vy = quaternion[bodyIndex].z;
			float vz = quaternion[bodyIndex].w;

			quaternion[bodyIndex].x = s*ds - vx*dvx - vy*dvy - vz*dvz;
			quaternion[bodyIndex].y = ds*vx + s*dvx + dvy*vz - dvz*vy;
			quaternion[bodyIndex].z = ds*vy + s*dvy + dvz*vx - dvx*vz;
			quaternion[bodyIndex].w = ds*vz + s*dvz + dvx*vy - dvy*vx;
		}
	}

	//Update rotation matrix
	{
		//Normalize quaternion
		{
			float mag2 = quaternion[bodyIndex].x*quaternion[bodyIndex].x +
				quaternion[bodyIndex].y*quaternion[bodyIndex].y +
				quaternion[bodyIndex].z*quaternion[bodyIndex].z +
				quaternion[bodyIndex].w*quaternion[bodyIndex].w;

			if (mag2!=0.0f && (fabs(mag2 - 1.0f) > 0.00001f)) {
				float mag = sqrt(mag2);
				quaternion[bodyIndex].x /= mag;
				quaternion[bodyIndex].y /= mag;
				quaternion[bodyIndex].z /= mag;
				quaternion[bodyIndex].w /= mag;
			}
		}
		float w = quaternion[bodyIndex].x;
		float x = quaternion[bodyIndex].y;
		float y = quaternion[bodyIndex].z;
		float z = quaternion[bodyIndex].w;

		float xx = x * x;
		float yy = y * y;
		float zz = z * z;
		float xy = x * y;
		float xz = x * z;
		float yz = y * z;
		float wx = w * x;
		float wy = w * y;
		float wz = w * z; 

		rotationMatrix[matrixIndex] = 1.0f-2.0f*(yy+zz);
		rotationMatrix[matrixIndex+1] = 2.0f*(xy-wz);
		rotationMatrix[matrixIndex+2] = 2.0f*(xz+wy);
		rotationMatrix[matrixIndex+3] = 2.0f*(xy+wz);
		rotationMatrix[matrixIndex+4] = 1.0f-2.0f*(xx+zz);
		rotationMatrix[matrixIndex+5] = 2.0f*(yz-wx);
		rotationMatrix[matrixIndex+6] = 2.0f*(xz-wy);
		rotationMatrix[matrixIndex+7] = 2.0f*(yz+wx);
		rotationMatrix[matrixIndex+8] = 1.0f-2.0f*(xx+yy);

	}

	//Update body VBO
	{
		
		float side = particleRadius * 3.0f;

		float normals[6][3] = { // Cube face normals 
			{ -1.0f, 0.0f, 0.0f }, { 0.0f, 1.0f, 0.0f }, { 1.0f, 0.0f, 0.0f },
			{ 0.0f, -1.0f, 0.0f }, { 0.0f, 0.0f, 1.0f }, { 0.0f, 0.0f, -1.0f } };
		int faceIndices[6][4] = { // Cube faces' vertex indices 
			{ 0, 1, 2, 3 }, { 3, 2, 6, 7 }, { 7, 6, 5, 4 },
			{ 4, 5, 1, 0 }, { 5, 6, 2, 1 }, { 7, 4, 0, 3 } };
		float vertices[8][3] = { // Cube vertex positions 
			{ -side, -side, side }, { -side, -side, -side }, { -side, side, -side },
			{ -side, side, side }, { side, -side, side }, { side, -side, -side },
			{ side, side, -side }, { side, side, side } };

		for (int i=0; i<6; i++) { //for every face
			for (int j=0; j<4; j++) { //for every vertex in the face

				float r0 = rotationMatrix[matrixIndex];
				float r1 = rotationMatrix[matrixIndex+1];
				float r2 = rotationMatrix[matrixIndex+2];
				float r3 = rotationMatrix[matrixIndex+3];
				float r4 = rotationMatrix[matrixIndex+4];
				float r5 = rotationMatrix[matrixIndex+5];
				float r6 = rotationMatrix[matrixIndex+6];
				float r7 = rotationMatrix[matrixIndex+7];
				float r8 = rotationMatrix[matrixIndex+8];

				float v0 = vertices[faceIndices[i][j]][0];
				float v1 = vertices[faceIndices[i][j]][1];
				float v2 = vertices[faceIndices[i][j]][2];
				
				bodyVBO[bodyVBOIndex] = r0*v0 + r1*v1 + r2*v2 + position[bodyIndex].x;
				bodyVBO[bodyVBOIndex+1] = r3*v0 + r4*v1 + r5*v2 + position[bodyIndex].y;
				bodyVBO[bodyVBOIndex+2] = r6*v0 + r7*v1 + r8*v2 + position[bodyIndex].z;

				float n0 = normals[i][0];
				float n1 = normals[i][1];
				float n2 = normals[i][2];

				bodyVBO[bodyVBOIndex+bodyVBOStride] = r0*n0 + r1*n1 + r2*n2;
				bodyVBO[bodyVBOIndex+bodyVBOStride+1] = r3*n0 + r4*n1 + r5*n2;
				bodyVBO[bodyVBOIndex+bodyVBOStride+2] = r6*n0 + r7*n1 + r8*n2;

				bodyVBOIndex += 3;
			}
		}
	}
}

__kernel
void updateParticles(
	__global float3* bodyPosition,
	__global float3* bodyVelocity,
	__global float* rotationMatrix,
	__global float3* angularVelocity,
	__global float* particleVBO,
	__global float3* particlePosition,
	__global float3* particleVelocity,
	float particleRadius
) {
	unsigned int particleIndex = get_global_id(0);
	unsigned int bodyIndex = particleIndex / 27;
	unsigned int matrixIndex = bodyIndex * 9;

	float3 originalRelativePos;
	//Calculate original relative position
	{
		int relativeIndex = particleIndex % 27;

		int xIndex = relativeIndex / 9;
		relativeIndex -= xIndex * 9;

		int yIndex = relativeIndex / 3;
		relativeIndex -= yIndex * 3;

		int zIndex = relativeIndex;

		float space = 2.0f*particleRadius;

		xIndex--;
		yIndex--;
		zIndex--;

		originalRelativePos.x = xIndex*space;
		originalRelativePos.y = yIndex*space;
		originalRelativePos.z = zIndex*space;
	}

	//Update particle position
	{

		particlePosition[particleIndex].x = 
			originalRelativePos.x*rotationMatrix[matrixIndex] + 
			originalRelativePos.y*rotationMatrix[matrixIndex+1] + 
			originalRelativePos.z*rotationMatrix[matrixIndex+2];

		particlePosition[particleIndex].y = 
			originalRelativePos.x*rotationMatrix[matrixIndex+3] + 
			originalRelativePos.y*rotationMatrix[matrixIndex+4] + 
			originalRelativePos.z*rotationMatrix[matrixIndex+5];

		particlePosition[particleIndex].z = 
			originalRelativePos.x*rotationMatrix[matrixIndex+6] + 
			originalRelativePos.y*rotationMatrix[matrixIndex+7] + 
			originalRelativePos.z*rotationMatrix[matrixIndex+8];

		particlePosition[particleIndex].x += bodyPosition[bodyIndex].x;
		particlePosition[particleIndex].y += bodyPosition[bodyIndex].y;
		particlePosition[particleIndex].z += bodyPosition[bodyIndex].z;

	}

	//Update particle velocity
	float scalar = sqrt(
			angularVelocity[bodyIndex].x*angularVelocity[bodyIndex].x + 
			angularVelocity[bodyIndex].y*angularVelocity[bodyIndex].y +
			angularVelocity[bodyIndex].z*angularVelocity[bodyIndex].z );

	scalar *= scalar;

	particleVelocity[particleIndex].x = bodyVelocity[bodyIndex].x;
	particleVelocity[particleIndex].y = bodyVelocity[bodyIndex].y;
	particleVelocity[particleIndex].z = bodyVelocity[bodyIndex].z;

	if (scalar > 0.0f) {
		float3 relativePosition = {
			particlePosition[particleIndex].x-bodyPosition[bodyIndex].x,
			particlePosition[particleIndex].y-bodyPosition[bodyIndex].y,
			particlePosition[particleIndex].z-bodyPosition[bodyIndex].z };

		float scalar2 = (
				angularVelocity[bodyIndex].x*relativePosition.x +
				angularVelocity[bodyIndex].y*relativePosition.y +
				angularVelocity[bodyIndex].z*relativePosition.z
				) / scalar;

		float3 term = {
			relativePosition.x - angularVelocity[bodyIndex].x*scalar2,
			relativePosition.y - angularVelocity[bodyIndex].y*scalar2,
			relativePosition.z - angularVelocity[bodyIndex].z*scalar2 };

		particleVelocity[particleIndex].x += (angularVelocity[bodyIndex].y*term.z - angularVelocity[bodyIndex].z*term.y);
		particleVelocity[particleIndex].y += (angularVelocity[bodyIndex].z*term.x - angularVelocity[bodyIndex].x*term.z);
		particleVelocity[particleIndex].z += (angularVelocity[bodyIndex].x*term.y - angularVelocity[bodyIndex].y*term.x);
	}

	//Update particle VBO
	unsigned int particleVBOIndex = particleIndex * 3;
	particleVBO[particleVBOIndex] = particlePosition[particleIndex].x;
	particleVBO[particleVBOIndex+1] = particlePosition[particleIndex].y;
	particleVBO[particleVBOIndex+2] = particlePosition[particleIndex].z;
}



